# -*- coding: utf-8 -*-
# backend/middleware/websocket_middleware.py
from fastapi import WebSocket
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response
import time
import json
from typing import Callable

from backend.utils.logger import logger


class WebSocketLoggingMiddleware(BaseHTTPMiddleware):
    """WebSocket日志中间件"""
    
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        start_time = time.time()
        
        # 记录请求信息
        logger.info(f"WebSocket请求 - 路径: {request.url.path}, 客户端: {request.client.host}")
        
        try:
            response = await call_next(request)
            
            # 记录响应信息
            process_time = time.time() - start_time
            logger.info(f"WebSocket响应 - 耗时: {process_time:.3f}s")
            
            return response
        
        except Exception as e:
            process_time = time.time() - start_time
            logger.error(f"WebSocket错误 - 耗时: {process_time:.3f}s, 错误: {str(e)}")
            raise


class WebSocketAuthMiddleware:
    """WebSocket认证中间件"""
    
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket
    
    async def authenticate(self, token: str) -> bool:
        """验证WebSocket连接的token"""
        try:
            # 这里可以添加token验证逻辑
            # 例如JWT验证、数据库查询等
            if not token:
                return False
            
            # 验证token格式和有效性
            # ... token验证逻辑
            
            return True
        
        except Exception as e:
            logger.error(f"WebSocket认证失败: {str(e)}")
            return False
